﻿from Game.G3Maze.env import MazeEnv
from Game.G3Maze.agent import MazeAgent


def train_agent(env, agent, steps=1000):
    total_reward = 0
    episode = 0

    while steps > 0:
        # 初始化环境
        current_state = env.get_current_position()["state_index"]
        episode_steps = 0

        while True:
            # 选择动作
            action = agent.choose_action()  # 0,1,2,3 # 行 列 左下右上

            # 执行动作
            _, reward, done, truncated, info = env.step(action)
            next_state = env.get_current_position()["state_index"]

            # 更新 Q 表
            agent.update(current_state, action, reward, next_state)

            # 统计累计奖励
            total_reward += reward
            current_state = next_state
            steps -= 1
            episode_steps += 1

            # 终止条件
            if done or truncated or steps <= 0:
                break

        print(
            f"Episode {episode} | Steps: {episode_steps} | Total Reward: {total_reward:.2f}"
        )
        episode += 1
        env.env.reset()


# 创建环境和智能体
env = MazeEnv()
agent = MazeAgent(env)
train_agent(env, agent)
